from rag import *
import pickle
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
# 使用索引检索
# 加载索引
with open(index_path(), 'rb') as f:
    index = pickle.load(f)

# 加载文档向量
with open(vector_base_path(), 'rb') as f:
    document_vectors = pickle.load(f)
# 用一个中文查询来检索文档
query_vector = vectorize([query])  # 将查询转换为向量

# 检索最相似的文档
D, I = index.search(query_vector, top_k)  # 返回距离D和索引I

documents=load_documents()
# 打印检索到的文档
for index,i in enumerate(I[0]): # 可以批量查询
    print(f"-----检索第 {index+1} 个结果---------- {0-D[0][index]}:\n{documents[i]}")
